import argparse
import random
import os


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('in_file', help='path to mxnet index')
    parser.add_argument('train', help='0~10', type=int)
    parser.add_argument('validation', help='0~10', type=int)
    parser.add_argument('test', help='0~10', type=int)

    return parser.parse_args()


def main():
    args = parse_args()
    assert(args.train + args.validation + args.test == 10)
    with open(args.in_file) as f:
        contents = []
        for line in f:
            contents.append(line)
        random.shuffle(contents)
        validation_num = int(float(len(contents)) / 10 * float(args.validation))
        test_num = int(float(len(contents)) / 10 * float(args.test))
        train_num = len(contents) - validation_num - test_num

        train_set = contents[0:train_num]
        validation_set = contents[train_num:train_num + validation_num]
        test_set = contents[train_num + validation_num:len(contents)]

        basename = os.path.splitext(os.path.basename(args.in_file))[0]
        if train_num != 0:
            with open(basename + '_train.txt', 'w') as out:
                print("We use balance here!!!")
                all_num = 0  # balance
                num = 0
                new_train_set = []
                for line in train_set:
                    info = line.split(' ')
                    if int(info[1]) == 2:
                        all_num += 1
                for line in train_set:
                    info = line.split(' ')
                    if num >= all_num and int(info[1]) == 1:
                        continue
                    if int(info[1]) == 1:
                        num += 1
                    new_train_set.append(line)
                random.shuffle(new_train_set)
                for line in new_train_set:
                    out.write(line)

        if validation_num != 0:
            with open(basename + '_validation.txt', 'w') as out:
                for line in validation_set:
                    out.write(line)

        if test_num != 0:
            with open(basename + '_test.txt', 'w') as out:
                for line in test_set:
                    out.write(line)

if __name__ == '__main__':
    main()
